import os
import time
import torch
from diffusers import StableDiffusionPipeline

def log(message):
    """打印带时间戳的日志"""
    from datetime import datetime
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print(f"[{timestamp}] {message}")

def main():
    start_time = time.time()
    log("开始处理...")

    try:
        # 设置模型路径
        base_model_path = 'E:/workspace/llm/text2image/models/AI-ModelScope/stable-diffusion-v1-5'

        # 检查模型文件
        log(f"检查模型文件夹: {base_model_path}")
        if not os.path.exists(base_model_path):
            raise FileNotFoundError(f"模型文件夹不存在: {base_model_path}")

        # 列出模型文件夹中的文件
        files = os.listdir(base_model_path)
        log(f"模型文件夹包含 {len(files)} 个文件")
        for file in files[:5]:  # 只显示前5个文件
            file_path = os.path.join(base_model_path, file)
            if os.path.isfile(file_path):
                log(f"- {file}: {os.path.getsize(file_path) / (1024*1024):.2f} MB")

        # 设置设备
        device = "cuda" if torch.cuda.is_available() else "cpu"
        log(f"使用设备: {device}")
        if device == "cuda":
            log(f"CUDA设备: {torch.cuda.get_device_name(0)}")
            log(f"CUDA可用内存: {torch.cuda.get_device_properties(0).total_memory / (1024**3):.2f} GB")

        # 加载基础模型
        log("加载基础模型...")
        pipe = StableDiffusionPipeline.from_pretrained(
            base_model_path,
            torch_dtype=torch.float32,  # 使用float32以提高兼容性
            safety_checker=None,
            requires_safety_checker=False
        ).to(device)

        # 生成图像
        log("生成图像中...")
        prompt = "A beautiful girl in the sunshine, high quality"

        # 使用简单设置
        with torch.no_grad():
            result = pipe(prompt, num_inference_steps=20)

        # 保存图像
        output_path = "result_simple2.png"
        result.images[0].save(output_path)
        log(f"图像已保存到: {output_path}")

        # 显示图像信息
        width, height = result.images[0].size
        log(f"图像尺寸: {width}x{height}")
        log(f"图像模式: {result.images[0].mode}")

        # 检查图像是否为全黑
        image_array = torch.tensor(list(result.images[0].getdata())).float()
        avg_pixel = image_array.mean().item()
        log(f"平均像素值: {avg_pixel:.2f} (0=黑色, 255=白色)")

        # 清理资源
        del pipe
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        # 计算总时间
        elapsed_time = time.time() - start_time
        log(f"处理完成! 总耗时: {elapsed_time:.2f}秒")

    except Exception as e:
        log(f"发生错误: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()